﻿using System;
using System.Collections.Generic;
using System.ServiceModel.Channels;
using System.Transactions;
using SimpleServiceBus.Endpoints.DotNetServiceBus.Transport;

namespace SimpleServiceBus.Endpoints.DotNetServiceBus.Utilities
{
    public class TransactionalQueueResourceManager : IEnlistmentNotification
    {
        private readonly IQueueClient _queue;
        private List<Message> _messages = new List<Message>();
        private Transaction _tx;
        private Exception _ex;

        public TransactionalQueueResourceManager(IQueueClient queue)
        {
            if (queue == null)
                throw new ArgumentNullException("queue");
            _queue = queue;
        }

        public Message Retrieve(TimeSpan timeout)
        {
            EnsureTransactionEnlistment();

            Message message;
            if (_tx != null)
            {
                try
                {
                    message = _queue.Peek(timeout);
                    _messages.Add(message);

                }
                catch (Exception ex)
                {
                    _ex = ex;
                    throw;
                }
            }
            else
            {
                message = _queue.Retrieve(timeout);
            }
            return message;

        }

        public IEnumerable<Message> RetrieveMultiple(int maxMessages,TimeSpan timeout)
        {
            EnsureTransactionEnlistment();

            IEnumerable<Message> messages;

            if (_tx != null)
            {
                try
                {
                    messages = _queue.PeekMultiple(maxMessages,timeout);
                    _messages.AddRange(messages);

                }
                catch (Exception ex)
                {
                    _ex = ex;
                    throw;
                }
            }
            else
            {
                messages = _queue.RetrieveMultiple(maxMessages,timeout);
            }
            return messages;
        }

        private void EnsureTransactionEnlistment()
        {
            if (_tx == null && Transaction.Current != null)
            {
                _tx = Transaction.Current;
                _tx.EnlistVolatile(this, EnlistmentOptions.None);
            }
        }

        void IEnlistmentNotification.Prepare(PreparingEnlistment preparingEnlistment)
        {
            if (_ex != null)
                preparingEnlistment.ForceRollback(_ex);

            var deleted = new List<Message>();

            try
            {
                foreach (Message message in _messages)
                {
                    if (message != null)
                    {
                        deleted.Add(message);
                        _queue.Delete(message);
                    }
                }
                preparingEnlistment.Prepared();
            }
            catch (Exception ex)
            {
                preparingEnlistment.ForceRollback(ex);
            }
            _messages = deleted;
            
        }

        void IEnlistmentNotification.Commit(Enlistment enlistment)
        {
            _messages.Clear();
            _tx = null;
        }

        void IEnlistmentNotification.Rollback(Enlistment enlistment)
        {
            _messages.ForEach(m => { if (m != null)_queue.Release(m); });

            _messages.Clear();
            _tx = null;
        }

        private void ReplaceDeletedMessages()
        {
            foreach (Message message in _messages)
                _queue.Send(message);
        }

        public void InDoubt(Enlistment enlistment)
        {
        }
    }
}
